#
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from apps.fmcw.fmcw_rsp import FmcwRsp
import apps.fmcw.t3 as t3

# Radar parameters
c = 3e8 #speed of light
BW = 150e6 #bandwidth
fc = 77e9 #carrier frequency
numADC = 256 # of adc samples
numChirps = 256 # of chirps per frame
numCPI = 10
T = 10e-6 # PRI 
PRF = 1/T
F = numADC/T # sampling frequency
dt = 1/F #sampling interval
slope = BW/T
lambda_ = c/fc # 波长
N = numChirps*numADC*numCPI # total # of adc samples
t = np.linspace(0, T * numChirps * numCPI, N) # linspace(0,T*numChirps*numCPI,N) # time axis, one frame
t_onePulse = np.arange(0, dt * numADC, dt)
numTX = 1
numRX = 8
Vmax = lambda_/(T*4) # Max Unamb velocity m/s
DFmax = 1/2*PRF # = Vmax/(c/fc/2); % Max Unamb Dopp Freq
dR = c/(2*BW) # range resol
Rmax = F*c/(2*slope) # TI's MIMO Radar doc
Rmax2 = c/2/PRF #lecture 2.3
dV = lambda_/(2*numChirps*T) # velocity resol, lambda/(2*framePeriod)
d_rx = lambda_/2 #dist. between rxs
d_tx = 4*d_rx # dist. between txs

N_Dopp = numChirps # length of doppler FFT
N_range = numADC #length of range FFT
N_azimuth = numTX*numRX
R = np.arange(0, Rmax, dR)  # range axis
V = np.linspace(-Vmax, Vmax, numChirps)  # Velocity axis
ang_ax = np.arange(-90, 91)  # angle axis



# Targets

r1_radial = 50
tar1_angle = -15
r1_y = np.cos(np.radians(tar1_angle)) * r1_radial
r1_x = np.sin(np.radians(tar1_angle)) * r1_radial
v1_radial = 10  # velocity 1
v1_y = np.cos(np.radians(tar1_angle)) * v1_radial
v1_x = np.sin(np.radians(tar1_angle)) * v1_radial
r1 = [r1_x, r1_y, 0]

r2_radial = 100
tar2_angle = 10
r2_y = np.cos(np.radians(tar2_angle)) * r2_radial
r2_x = np.sin(np.radians(tar2_angle)) * r2_radial
v2_radial = -15  # velocity 2
v2_y = np.cos(np.radians(tar2_angle)) * v2_radial
v2_x = np.sin(np.radians(tar2_angle)) * v2_radial
r2 = [r2_x, r2_y, 0]

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

tx_loc = np.zeros((numTX, 3))
for i in range(numTX):
    tx_loc[i][0], tx_loc[i][1], tx_loc[i][2] = i * d_tx, 0, 0
    ax.scatter(tx_loc[i][0], tx_loc[i][1], tx_loc[i][2], color='blue', marker='o')
rx_loc = np.zeros((numRX, 3))
for i in range(numRX):
    rx_loc[i][0], rx_loc[i][1], rx_loc[i][2] = tx_loc[numTX-1][0] + d_tx + i * d_rx, 0, 0
    ax.scatter(rx_loc[i][0], rx_loc[i][1], rx_loc[i][2], color='red', marker='o')

ax.set_xlabel('X-axis')
ax.set_ylabel('Y-axis')
ax.set_zlabel('Z-axis')
ax.set_title('Transmitter and Receiver Locations')
plt.show()

tar1_loc = np.zeros((len(t), 3))
tar2_loc = np.zeros((len(t), 3))
tar1_loc[:, 0] = r1[0] + v1_x * t
tar2_loc[:, 0] = r2[0] + v2_x * t
tar1_loc[:, 1] = r1[1] + v1_y * t
tar2_loc[:, 1] = r2[1] + v2_y * t

# TX
delays_tar1 = np.zeros((numTX, numRX, len(t)))
delays_tar2 = np.zeros((numTX, numRX, len(t)))

for i in range(numTX):
    for j in range(numRX):
        delays_tar1[i, j] = (np.linalg.norm(tar1_loc - rx_loc[j], ord=2, axis=1) + np.linalg.norm(tar1_loc - tx_loc[i], ord=2, axis=1))/c
        delays_tar2[i, j] = (np.linalg.norm(tar2_loc - rx_loc[j], ord=2, axis=1) + np.linalg.norm(tar2_loc - tx_loc[i], ord=2, axis=1))/c

# Complex signal
def phase(tx, fx, slope):
    return 2 * np.pi * (fx * tx + slope / 2 * tx ** 2)

def phase2(tx, fx, r, v):
    return 2 * np.pi * (2 * fx * r / c + tx * (2 * fx * v / c + 2 * slope * r / c))

fr1 = 2 * r1[1] * slope / c
fr2 = 2 * r2[1] * slope / c
fd1 = 2 * v1_radial * fc / c  # doppler freq
fd2 = 2 * v2_radial * fc / c
f_if1 = fr1 + fd1  # beat or IF freq
f_if2 = fr2 + fd2

signal_t = np.zeros(N, dtype=complex)
signal_1 = np.zeros(N, dtype=complex)
signal_2 = np.zeros(N, dtype=complex)
mixed = np.zeros((numTX, numRX, N), dtype=complex)

for i in range(numTX):
    for j in range(numRX):
        print(f'Processing Channel: {j + 1}/{numRX}')
        for k in range(numChirps * numCPI):
            phase_t = phase(t_onePulse, fc, slope)
            phase_1 = phase(t_onePulse - delays_tar1[i, j, k * numADC:(k + 1) * numADC], fc, slope)  # received
            phase_2 = phase(t_onePulse - delays_tar2[i, j, k * numADC:(k + 1) * numADC], fc, slope)
            signal_t[k * numADC:(k + 1) * numADC] = np.exp(1j * phase_t)
            signal_1[k * numADC:(k + 1) * numADC] = np.exp(1j * (phase_t - phase_1))
            signal_2[k * numADC:(k + 1) * numADC] = np.exp(1j * (phase_t - phase_2))
        mixed[i, j] = signal_1 + signal_2

# Plotting
plt.figure()
plt.subplot(3, 1, 1)
plt.plot(t, signal_t.real)
plt.title('TX')
plt.xlim([0, 0.1e-4])
plt.xlabel('Time (sec)')
plt.ylabel('Amplitude')

plt.subplot(3, 1, 2)
plt.plot(t, signal_1.real)
plt.title('RX')
plt.xlim([0, 0.1e-4])
plt.xlabel('Time (sec)')
plt.ylabel('Amplitude')

plt.subplot(3, 1, 3)
plt.plot(t, mixed[0][7].real)
plt.title('Mixed')
plt.xlim([0, 0.1e-4])
plt.xlabel('Time (sec)')
plt.ylabel('Amplitude')

plt.tight_layout()
plt.show()

# Initialize RDMs array
mixed_t = mixed.squeeze(0).transpose(1,0)
RDC = None
for i in range(numADC):
    ddi = mixed_t[i::numADC]
    if RDC is None:
        RDC = np.array([ddi])
    else:
        RDC = np.vstack((RDC, np.array([ddi])))
RDMs = np.zeros((N_range, N_Dopp, numTX * numRX, numCPI), dtype=complex)
# Perform 2-D FFT on each CPI frame
for i in range(numCPI):
    RD_frame = RDC[:, i * numChirps:(i + 1) * numChirps, :] # CPI一帧的数据
    fft2_result = None
    for ridx in range(numTX*numRX):
        fri = np.fft.fft2(RD_frame[:,:,ridx])
        if fft2_result is None:
            fft2_result = np.array([fri]).transpose(1,2,0)
        else:
            fft2_result = np.concatenate((fft2_result, np.array([fri]).transpose(1,2,0)), axis=-1)
    fftshifted_result = None
    for ridx in range(numTX*numRX):
        fsri = np.fft.fftshift(fft2_result[:,:,ridx], axes=(-1,))
        if fftshifted_result is None:
            fftshifted_result = np.array([fsri]).transpose(1,2,0)
        else:
            fftshifted_result = np.concatenate((fftshifted_result, np.array([fsri]).transpose(1,2,0)), axis=-1)
    RDMs[:, :, :, i] = fftshifted_result

# Plotting
plt.figure()
# Adjust imshow to match the shape of RDMs
img_data = 20 * np.log10(np.abs(RDMs[:, :, 0, 0]) / np.max(np.max(np.abs(RDMs[:, :, 0, 0]))))
img_data = img_data[::-1] # 将图像行进行倒排，第1行变为最后一行，依此类推
img0 = plt.imshow(img_data, extent=[V[0], V[-1], R[0], R[-1]], aspect='auto')
# 获取当前图像的颜色限制
clim = img0.get_clim()
# 设置颜色轴的限制
plt.clim(clim[0]/2, 0)
plt.xlabel('Velocity (m/s)')
plt.ylabel('Range (m)')
plt.show()


# CA-CFAR
numGuard = 2 # of guard cells
numTrain = numGuard*2 # of training cells
P_fa = 1e-5 # desired false alarm rate
SNR_OFFSET = -5 # dB
RDM_dB = 10*np.log10(abs(RDMs[:,:,1,1])/np.max(abs(RDMs[:,:,1,1])))
RDM_mask, cfar_ranges, cfar_dopps, K = FmcwRsp.ca_cfar(RDM_dB, numGuard, numTrain, P_fa, SNR_OFFSET)
# 绘制CAFR图像 imagesc(V,R,20*log10(abs(RDMs(:,:,1,1))/max(max(abs(RDMs(:,:,1,1))))));
plt.figure()
img_data = RDM_mask[::-1]
img1 = plt.imshow(img_data, extent=[V[0], V[-1], R[0], R[-1]], aspect='auto')
plt.xlabel('Velocity (m/s)')
plt.ylabel('Range (m)')
plt.title('CA-CFAR')
plt.show()


# Angle Estimation - MUSIC Pseudo Spectrum
d = 0.5
M = numCPI # of snapshots
# 初始化a1数组，假设其形状为[numTX*numRX, len(ang_ax)]
a1 = np.zeros((numTX * numRX, len(ang_ax)), dtype=np.complex_)
# 计算a1
for k in range(len(ang_ax)):
    # 创建一个从0到numTX*numRX-1的列向量
    idx = np.arange(numTX * numRX)
    # 计算复指数并赋值给a1的第k列
    a1[:, k] = np.exp(-1j * 2 * np.pi * d * (idx * np.sin(np.radians(ang_ax[k]))))
music_spectrum = np.zeros((K, len(ang_ax)))
for i in range(K):
    Rxx = np.zeros((numTX * numRX, numTX * numRX))
    for m in range(M):
        # A = np.squeeze(RDMs[cfar_ranges[i], cfar_dopps[i], :, m])
        A = RDMs[cfar_ranges[i], cfar_dopps[i], :, m]
        A = np.expand_dims(A, axis=1)
        Rxx = Rxx + 1/M * np.dot(A, A.conj().T)
    D, Q = np.linalg.eig(Rxx)  # Q: eigenvectors (columns), D: eigenvalues
    I = np.argsort(D.real)[::-1]  # Sort eigenvalues in descending order
    Q = Q[:, I]  # Sort the eigenvectors accordingly
    Qs = Q[:, 0]  # Get the signal eigenvector
    Qn = Q[:, 1:]  # Get the noise eigenvectors
    for k in range(len(ang_ax)):
        # music_spectrum(i,k)=(a1(:,k)'*a1(:,k))/(a1(:,k)'*(Qn*Qn')*a1(:,k));
        # 1*8 @ 8*1      1*8 @ (8*7 @ 7*8) => 1*8 @ 8*1
        a1k = np.expand_dims(a1[:, k], 1)
        music_spectrum[i, k] = np.real(np.dot(a1k.conj().T, a1k) / np.dot(np.dot(a1k.conj().T, np.dot(Qn, Qn.conj().T)), a1k))
        # music_spectrum[i, k] = np.matmul(a1[:,k].T, a1[:,k]) / (   np.matmul(np.matmul(a1[:,k].T, np.matmul(Qn, Qn.T)), a1[:,k]) )
# Plotting the MUSIC spectrum
plt.figure()
# plt.hold(True)
plt.grid(True)
plt.title('MUSIC Spectrum')
plt.xlabel('Angle in degrees')
for k in range(K):
    plt.plot(ang_ax, np.log10(np.abs(music_spectrum[k, :])))
plt.show()    


# 使用 NumPy 的 FFT 函数代替 MATLAB 的 fft
rangeFFT = np.fft.fft(RDC, axis=0)
range_az_music = np.zeros((N_range, len(ang_ax)))

for i in range(N_range):
    Rxx = np.zeros((numTX * numRX, numTX * numRX), dtype=np.complex_)
    for m in range(M):
        # MATLAB 中的 squeeze 和 sum 函数对应于 NumPy 的 squeeze 和 sum 函数
        A = np.sum(rangeFFT[i, m * numChirps:(m+1) * numChirps, :], axis=0).reshape(-1, 1)
        Rxx += (1 / M) * np.dot(A, A.conj().T)
    D, Q = np.linalg.eig(Rxx)  # Q: eigenvectors (columns), D: eigenvalues
    I = np.argsort(D.real)[::-1]  # Sort eigenvalues in descending order
    Q = Q[:, I]  # Sort the eigenvectors accordingly
    Qs = Q[:, 0]  # Get the signal eigenvector
    Qn = Q[:, 1:]  # Get the noise eigenvectors
    music_spectrum2 = np.zeros(len(ang_ax))
    for k in range(len(ang_ax)):
        # NumPy 中的矩阵乘法使用 np.dot
        music_spectrum2[k] = np.real(np.dot(a1[:, k].conj().T, a1[:, k]) / np.dot(np.dot(a1[:, k].conj().T, np.dot(Qn, Qn.conj().T)), a1[:, k]))
    range_az_music[i, :] = music_spectrum2

# 绘制图像
plt.figure()
img_data = 20 * np.log10(np.abs(range_az_music) / np.max(np.abs(range_az_music)))
img_data = img_data[::-1]
plt.imshow(img_data, 
           extent=[ang_ax[0], ang_ax[-1], R[0], R[-1]], 
           aspect='auto', 
           cmap='jet')
plt.colorbar()
plt.xlabel('Azimuth')
plt.ylabel('Range (m)')
plt.title('MUSIC Range-Angle Map')
plt.clim(vmin=np.min(20 * np.log10(np.abs(range_az_music) / np.max(np.abs(range_az_music)))), 
         vmax=np.max(20 * np.log10(np.abs(range_az_music) / np.max(np.abs(range_az_music)))))
plt.show()

# 生成点云
# 找到music_spectrum每行的最大值索引
I1 = np.argmax(music_spectrum[1, :])
angle1 = ang_ax[I1]
I2 = np.argmax(music_spectrum[0, :])
angle2 = ang_ax[I2]
# 计算坐标
coor1 = np.array([cfar_ranges[1] * np.cos(np.deg2rad(angle1)), 
                  cfar_ranges[1] * np.sin(np.deg2rad(angle1)), 0])
coor2 = np.array([cfar_ranges[0] * np.cos(np.deg2rad(angle2)), 
                  cfar_ranges[0] * np.sin(np.deg2rad(angle2)), 0])
# 创建3D散点图
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(coor1[0], coor1[1], coor1[2], s=10, c='m', marker='o', edgecolors='k', linewidth=3)
ax.scatter(coor2[0], coor2[1], coor2[2], s=10, c='b', marker='o', edgecolors='k', linewidth=3)
# 设置图形标题和坐标轴标签
ax.set_title('3D Coordinates (Point Cloud) of the targets')
ax.set_xlabel('Range (m) X')
ax.set_ylabel('Range (m) Y')
ax.set_zlabel('Range (m) Z')
plt.show()
print(f'目标１： {coor1[0]}, {coor1[1]}, {coor1[2]}')
print(f'目标１： {coor2[0]}, {coor2[1]}, {coor2[2]}')
# 绘制２维散点图
plt.figure()
plt.scatter(coor1[0], coor1[1], s=10, c='m', marker='o', edgecolors='k', linewidth=3)
plt.scatter(coor2[0], coor2[1], s=10, c='b', marker='o', edgecolors='k', linewidth=3)
# 设置图形标题和坐标轴标签
plt.title('2D Coordinates (Point Cloud) of the targets')
plt.xlabel('Range (m) X')
plt.ylabel('Range (m) Y')
plt.show()